use crate::error::{Result, WsError};
use crate::protocol::{ConnectionInfo, Message};
use dashmap::DashMap;
use std::sync::Arc;
use tokio::sync::mpsc;
use tracing::{debug, error, info, warn};
use tungstenite::Message as WsMessage;
use uuid::Uuid;

pub type ConnectionMap = Arc<DashMap<Uuid, Connection>>;
pub type UserToConnectionMap = Arc<DashMap<Uuid, Uuid>>; // user_id -> connection_id

#[derive(Debug)]
pub struct Connection {
    pub id: Uuid,
    pub username: String,
    pub sender: mpsc::UnboundedSender<WsMessage>,
    pub connected_at: u64,
    pub last_activity: std::sync::atomic::AtomicU64,
}

impl Connection {
    pub fn new(
        id: Uuid,
        username: String,
        sender: mpsc::UnboundedSender<WsMessage>,
    ) -> Self {
        let now = current_timestamp();
        Self {
            id,
            username,
            sender,
            connected_at: now,
            last_activity: std::sync::atomic::AtomicU64::new(now),
        }
    }
    
    pub fn update_activity(&self) {
        let now = current_timestamp();
        self.last_activity.store(now, std::sync::atomic::Ordering::Relaxed);
    }
    
    pub fn get_info(&self) -> ConnectionInfo {
        ConnectionInfo {
            id: self.id,
            username: self.username.clone(),
            connected_at: self.connected_at,
            last_activity: self.last_activity.load(std::sync::atomic::Ordering::Relaxed),
        }
    }
    
    pub async fn send_message(&self, message: WsMessage) -> Result<()> {
        self.sender
            .send(message)
            .map_err(|e| WsError::SendError(e.to_string()))?;
        self.update_activity();
        Ok(())
    }
}

pub struct ConnectionManager {
    connections: ConnectionMap,
    user_to_connection: UserToConnectionMap,
    stats: Arc<ConnectionStats>,
}

#[derive(Debug, Default)]
pub struct ConnectionStats {
    pub total_connections: std::sync::atomic::AtomicU64,
    pub active_connections: std::sync::atomic::AtomicU64,
    pub messages_sent: std::sync::atomic::AtomicU64,
    pub messages_received: std::sync::atomic::AtomicU64,
}

impl ConnectionManager {
    pub fn new() -> Self {
        Self {
            connections: Arc::new(DashMap::new()),
            user_to_connection: Arc::new(DashMap::new()),
            stats: Arc::new(ConnectionStats::default()),
        }
    }
    
    pub fn get_connections(&self) -> &ConnectionMap {
        &self.connections
    }
    
    pub fn get_stats(&self) -> Arc<ConnectionStats> {
        Arc::clone(&self.stats)
    }
    
    pub fn add_connection(&self, connection: Connection, user_id: Uuid) {
        let connection_id = connection.id;
        self.connections.insert(connection_id, connection);
        self.user_to_connection.insert(user_id, connection_id);
        self.stats.active_connections.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        self.stats.total_connections.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }
    
    pub fn remove_connection(&self, connection_id: &Uuid) -> Option<Connection> {
        let result = self.connections.remove(connection_id).map(|(_, conn)| conn);
        if let Some(ref conn) = result {
            // Remove from user_to_connection mapping
            let user_ids_to_remove: Vec<Uuid> = self.user_to_connection
                .iter()
                .filter(|entry| *entry.value() == *connection_id)
                .map(|entry| *entry.key())
                .collect();
            
            for user_id in user_ids_to_remove {
                self.user_to_connection.remove(&user_id);
            }
            
            self.stats.active_connections.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
        }
        result
    }
    
    pub fn get_connection(&self, id: &Uuid) -> Option<dashmap::mapref::one::Ref<Uuid, Connection>> {
        self.connections.get(id)
    }
    
    pub fn get_all_connections(&self) -> Vec<ConnectionInfo> {
        self.connections
            .iter()
            .map(|entry| entry.value().get_info())
            .collect()
    }
    
    pub fn get_user_id_by_connection_id(&self, connection_id: &Uuid) -> Option<Uuid> {
        for entry in self.user_to_connection.iter() {
            if *entry.value() == *connection_id {
                return Some(*entry.key());
            }
        }
        None
    }
    
    pub async fn broadcast_message(&self, message: &Message, exclude_id: Option<Uuid>) -> Result<()> {
        let json_message = message.to_json().map_err(|e| WsError::SendError(e.to_string()))?;
        let ws_message = WsMessage::Text(json_message);
        
        let mut failed_connections = Vec::new();
        
        for entry in self.connections.iter() {
            let connection_id = *entry.key();
            if let Some(exclude) = exclude_id {
                if connection_id == exclude {
                    continue;
                }
            }
            
            if let Err(e) = entry.value().send_message(ws_message.clone()).await {
                error!("Failed to send message to connection {}: {}", connection_id, e);
                failed_connections.push(connection_id);
            } else {
                self.stats.messages_sent.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            }
        }
        
        // Remove failed connections
        for id in failed_connections {
            self.remove_connection(&id);
        }
        
        Ok(())
    }
    
    pub async fn send_direct_message(&self, target_user_id: Uuid, message: &Message) -> Result<()> {
        // First find the connection ID for this user
        if let Some(connection_id_entry) = self.user_to_connection.get(&target_user_id) {
            let connection_id = *connection_id_entry.value();
            if let Some(connection) = self.get_connection(&connection_id) {
                let json_message = message.to_json().map_err(|e| WsError::SendError(e.to_string()))?;
                let ws_message = WsMessage::Text(json_message);
                
                connection.send_message(ws_message).await?;
                self.stats.messages_sent.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                Ok(())
            } else {
                Err(WsError::ConnectionNotFound(target_user_id.to_string()))
            }
        } else {
            Err(WsError::ConnectionNotFound(target_user_id.to_string()))
        }
    }
    
    pub fn increment_received(&self) {
        self.stats.messages_received.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }
}

fn current_timestamp() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_secs()
}